from collections import defaultdict, deque
import copy
import datetime
import logging

from coin.proto.coin_strategy_summary_pb2 import (StrategySummary,
                                                  StrategyDetailEntry,
                                                  StrategyDetailSummary)
from coin.base.query_util import (
    query_exchange_rates)


def _convert_pnl_adj(pnl_adjs):
  pnl_adjs = copy.deepcopy(pnl_adjs)
  adj_dict = defaultdict(dict)
  assert isinstance(pnl_adjs, list)
  for adj in pnl_adjs:
    strategy_name = adj['strategy_name']
    accounting_currency = adj['accounting_currency']
    entry = adj_dict[strategy_name]
    if accounting_currency not in entry:
      for key in ('pnl_adj', 'otc_pnl_adj'):
        if key not in adj:
          adj[key] = 0.
      entry[accounting_currency] = adj
    else:
      entry[accounting_currency]['pnl_adj'] += adj.get('pnl_adj', 0)
      entry[accounting_currency]['otc_pnl_adj'] += adj.get('otc_pnl_adj', 0)
  return adj_dict


def _get_summary_after_pnl_adj(trading_date, pnl_adj_dict, summary_adj):
  summary_after_adj = []
  for strategy_name, summaries in summary_adj.items():
    assert len(summaries) > 0
    pnl_adj = pnl_adj_dict[strategy_name]
    accounting_currencies = pnl_adj.keys()
    for accounting_currency in accounting_currencies:
      candidates = [
          summary for summary in summaries if accounting_currency in
          summary.stat_per_quote
      ]
      if len(candidates) == 0:
        continue
      key_func = lambda x: x.stat_per_quote[accounting_currency].pnl_mark
      trading_adj_value = pnl_adj[accounting_currency].get('pnl_adj', 0)
      otc_adj_value = pnl_adj[accounting_currency].get('otc_pnl_adj', 0)
      total_adj_value = trading_adj_value + otc_adj_value
      if total_adj_value >= 0:
        target = min(candidates, key=key_func)
      else:
        target = max(candidates, key=key_func)
      stat = target.stat_per_quote[accounting_currency]
      if stat.pnl_mark == 0:
        logging.info(f'pnl mark = 0: {strategy_name}')
        price_in_usd = query_exchange_rates([accounting_currency], 'USD', trading_date)[0]
      else:
        price_in_usd = stat.pnl_mark_in_usd / stat.pnl_mark
      trading_adj_value_in_usd = price_in_usd * trading_adj_value
      otc_adj_value_in_usd = price_in_usd * otc_adj_value
      total_adj_value_in_usd = price_in_usd * total_adj_value
      stat.pnl_mark += total_adj_value
      if stat.HasField('pnl_mark_in_usd'):
        stat.pnl_mark_in_usd += total_adj_value_in_usd
      if stat.HasField('trading_pnl_mark'):
        stat.trading_pnl_mark += trading_adj_value
      if stat.HasField('trading_pnl_mark_in_usd'):
        stat.trading_pnl_mark_in_usd += trading_adj_value_in_usd
      if stat.HasField('pnl_mark_from_balance'):
        stat.pnl_mark_from_balance += trading_adj_value
      if stat.HasField('pnl_mark_from_balance_in_usd'):
        stat.pnl_mark_from_balance_in_usd += trading_adj_value_in_usd
      if stat.HasField('active_trading_pnl_mark'):
        stat.active_trading_pnl_mark += trading_adj_value
      if stat.HasField('active_trading_pnl_mark_in_usd'):
        stat.active_trading_pnl_mark_in_usd += trading_adj_value_in_usd
      if stat.HasField('otc_pnl_mark'):
        stat.otc_pnl_mark += otc_adj_value
      if stat.HasField('otc_pnl_mark_in_usd'):
        stat.otc_pnl_mark_in_usd += otc_adj_value_in_usd
    summary_after_adj += summaries
  return summary_after_adj


def _get_detail_after_pnl_adj(pnl_adj_dict, detail_adj):
  detail_after_adj = []
  for strategy_name, details in detail_adj.items():
    assert len(details) > 0
    pnl_adj = pnl_adj_dict[strategy_name]
    accounting_currencies = pnl_adj.keys()
    for accounting_currency in accounting_currencies:
      candidates = []
      for detail in details:
        for entry in detail.each_detail:
          if accounting_currency == entry.accounting_currency and \
             entry.type == StrategyDetailEntry.STRATEGY and \
             len(entry.pnl) > 0:
            candidates.append(entry)
      if len(candidates) == 0:
        continue
      key_func = lambda x: x.pnl[-1].pnl_mark_from_balance - \
                           x.pnl[0].pnl_mark_from_balance
      adj_value = pnl_adj[accounting_currency]['pnl_adj']
      if adj_value >= 0:
        target = min(candidates, key=key_func)
      else:
        target = max(candidates, key=key_func)
      pnl_raw = [stat.pnl_mark_from_balance for stat in target.pnl if
                 stat.HasField('pnl_mark_from_balance')]
      pnl_shifted = deque(pnl_raw)
      pnl_shifted.rotate(1)
      pnl_shifted[0] = 0
      pnl_diff = [raw - shifted for raw, shifted in zip(pnl_raw, pnl_shifted)]
      if adj_value >= 0:
        pnl_target_idx = pnl_diff.index(min(pnl_diff))
      else:
        pnl_target_idx = pnl_diff.index(max(pnl_diff))
      for stat in target.pnl[pnl_target_idx:]:
        stat.pnl_mark_from_balance += adj_value
    detail_after_adj += details
  return detail_after_adj


def _adjust_strat_summary_pnl(trading_date, summaries, pnl_adj_dict):
  if summaries is None:
    return summaries
  adj_strats = pnl_adj_dict.keys()
  summary_after_adj = []
  summary_adj = defaultdict(list)
  for summary in summaries:
    assert isinstance(summary, StrategySummary)
    strategy_name = summary.strategy.strategy_name
    if strategy_name not in adj_strats:
      summary_after_adj.append(summary)
      continue
    summary_copy = StrategySummary()
    summary_copy.CopyFrom(summary)
    summary_adj[strategy_name].append(summary_copy)
  summary_after_adj += _get_summary_after_pnl_adj(
      trading_date, pnl_adj_dict, summary_adj)
  return summary_after_adj


def _adjust_strat_detail_pnl(details, pnl_adj_dict):
  if details is None:
    return details
  adj_strats = pnl_adj_dict.keys()
  detail_after_adj = []
  detail_adj = defaultdict(list)
  for detail in details:
    assert isinstance(detail, StrategyDetailSummary)
    strategy_name = detail.strategy.strategy_name
    if strategy_name not in adj_strats:
      detail_after_adj.append(detail)
      continue
    detail_copy = StrategyDetailSummary()
    detail_copy.CopyFrom(detail)
    detail_adj[strategy_name].append(detail_copy)
  detail_after_adj += _get_detail_after_pnl_adj(pnl_adj_dict, detail_adj)
  return detail_after_adj


def _adjust_daily_pnl(summaries, details, trading_date, daily_pnl_adj):
  td_str = trading_date.strftime('%Y%m%d')
  if daily_pnl_adj is None or td_str not in daily_pnl_adj:
    summary_after_adj = summaries
    detail_after_adj = details
  else:
    pnl_adj_dict = _convert_pnl_adj(daily_pnl_adj[td_str])
    summary_after_adj = _adjust_strat_summary_pnl(trading_date, summaries, pnl_adj_dict)
    detail_after_adj = _adjust_strat_detail_pnl(details, pnl_adj_dict)
  return summary_after_adj, detail_after_adj


def adjust_pnl_from_config(summaries, details, start_date, end_date, daily_pnl_adj):
  assert end_date >= start_date, (start_date, end_date)
  trading_date = start_date
  while end_date >= trading_date:
    summaries, details = _adjust_daily_pnl(
        summaries, details, trading_date, daily_pnl_adj)
    summaries = [adjust_pnl_from_balance_inferred_in_usd(summary) for
                 summary in summaries]
    trading_date += datetime.timedelta(days=1)
  return summaries, details


def adjust_pnl_from_balance_inferred_in_usd(summary):
  assert isinstance(summary, StrategySummary)
  ccy_pnl_map = {
      ccy: stat.pnl_mark for ccy, stat in summary.stat_per_quote.items()}
  ccy_pnl_in_usd_map = {
      ccy: stat.pnl_mark_in_usd for ccy, stat in summary.stat_per_quote.items()}
  ccy_symbol_num_map = defaultdict(int)
  for _, acct_stat in summary.stat_per_account.items():
    for _, symbol_stat in acct_stat.stat_per_symbol.items():
      assert symbol_stat.HasField('accounting_currency'), summary
      ccy_symbol_num_map[symbol_stat.accounting_currency] += 1
  assert len(ccy_pnl_map) >= len(ccy_symbol_num_map)
  ccy_diff = set(list(ccy_pnl_map.keys())).difference(
      set(list(ccy_symbol_num_map.keys())))
  total_untracked_pnl_in_usd = sum([ccy_pnl_in_usd_map[ccy] for ccy in ccy_diff])
  total_symbol_num = sum(list(ccy_symbol_num_map.values()))
  untracked_pnl_in_usd_per_symbol = 0
  if total_untracked_pnl_in_usd != 0:
    # TODO(chensili): uncomment assertion after PIT product info is supported
    # in rare case, total symbol num may be zero if products are all delisted in previous
    # trading day.
    # assert total_symbol_num > 0, summary
    if total_symbol_num > 0:
      untracked_pnl_in_usd_per_symbol = total_untracked_pnl_in_usd / total_symbol_num
  summary_copy = StrategySummary()
  summary_copy.CopyFrom(summary)
  for _, acct_stat in summary_copy.stat_per_account.items():
    for _, symbol_stat in acct_stat.stat_per_symbol.items():
      ccy = symbol_stat.accounting_currency
      ccy_num = ccy_symbol_num_map[ccy]
      symbol_stat.pnl_mark_from_balance_inferred = ccy_pnl_map[ccy] / ccy_num
      symbol_stat.pnl_mark_from_balance_inferred_in_usd = \
          ccy_pnl_in_usd_map[ccy] / ccy_num + untracked_pnl_in_usd_per_symbol
  return summary_copy
